from sqlalchemy import Column, Integer, Text, DateTime, ForeignKey, Boolean, String, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship
from datetime import datetime
from sqlalchemy import UniqueConstraint
# 创建基础模型类
Base = declarative_base()


# 数据库模型定义
class Conversation(Base):
    __tablename__ = "conversations"

    conversation_id = Column(Integer, primary_key=True, autoincrement=True)
    title = Column(Text, nullable=False)
    start_time = Column(DateTime, nullable=False, default=datetime.utcnow)
    last_update_time = Column(DateTime, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)
    model_name = Column(Text)
    system_prompt_id = Column(Integer, nullable=True)
    kb_id = Column(Integer, nullable=True)  # 保存会话使用的知识库ID
    user_id = Column(Integer, ForeignKey("users.user_id", ondelete="SET NULL"), nullable=True)

    # 关联关系
    messages = relationship("Message", back_populates="conversation", cascade="all, delete-orphan")
    user = relationship("User", back_populates="conversations")


class Message(Base):
    __tablename__ = "messages"

    message_id = Column(Integer, primary_key=True, autoincrement=True)
    conversation_id = Column(Integer, ForeignKey("conversations.conversation_id", ondelete="CASCADE"), nullable=False)
    sender_type = Column(Integer, nullable=False)  # 0=用户，1=AI
    content = Column(Text, nullable=False)
    send_time = Column(DateTime, nullable=False, default=datetime.utcnow)
    sequence = Column(Integer, nullable=False)
    citations = Column(JSON, nullable=True)

    # 关联关系
    conversation = relationship("Conversation", back_populates="messages")

class SystemPrompt(Base):
    __tablename__ = "system_prompts"

    system_prompt_id = Column(Integer, primary_key=True, autoincrement=True)
    title = Column(String(255), nullable=True)
    content = Column(Text, nullable=False)
    enable_context = Column(Boolean, nullable=False, default=False)
    context_placeholder = Column(String(255), nullable=True)
    user_id = Column(Integer, ForeignKey("users.user_id", ondelete="SET NULL"), nullable=True)
    is_public = Column(Boolean, nullable=False, default=False)
    created_time = Column(DateTime, nullable=False, default=datetime.utcnow)
    updated_time = Column(DateTime, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)

    # 关联关系
    user = relationship("User", back_populates="system_prompts")

class Model(Base):
    __tablename__ = "models"
    
    model_id = Column(Integer, primary_key=True, autoincrement=True)
    model_name = Column(String(255), nullable=False, unique=True)
    display_name = Column(String(255), nullable=False)
    description = Column(Text)
    model_type = Column(String(50), nullable=False, default="chat")  # chat 或 embedding
    is_active = Column(Boolean, nullable=False, default=True)
    created_time = Column(DateTime, nullable=False, default=datetime.utcnow)
    updated_time = Column(DateTime, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)


class KnowledgeBase(Base):
    """知识库表，存储知识库基本信息"""
    __tablename__ = 'knowledge_bases'

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String(255), nullable=False, unique=True)
    description = Column(Text)
    embedding_model_name = Column(String(255), nullable=True)
    user_id = Column(Integer, ForeignKey("users.user_id", ondelete="SET NULL"), nullable=True)
    is_public = Column(Boolean, nullable=False, default=False)
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

    # 关联关系
    user = relationship("User", back_populates="knowledge_bases")


class Document(Base):
    """文档表，存储文档基本信息"""
    __tablename__ = 'documents'

    id = Column(Integer, primary_key=True, autoincrement=True)
    kb_id = Column(Integer, ForeignKey("knowledge_bases.id", ondelete="SET NULL"), nullable=True)
    doc_name = Column(String(255), nullable=False)
    doc_hash = Column(String(64), nullable=False)
    embedding_model_name = Column(String(255), nullable=True)
    chunk_mode = Column(String(50), nullable=True)  # paragraph/sentence/length/delimiter
    chunk_params = Column(JSON, nullable=True)     # 记录分块参数，如 {"delimiter":"##","chunk_size":500}
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    is_enabled = Column(Boolean, nullable=False, default=True)

    __table_args__ = (
        UniqueConstraint('kb_id', 'doc_hash', name='uq_kb_doc_hash'),
    )

    def __repr__(self):
        return f"<Document(doc_name='{self.doc_name}', doc_hash='{self.doc_hash}', kb_id={self.kb_id})>"


class VectorEmbedding(Base):
    """向量嵌入表，存储文本块和对应的向量"""
    __tablename__ = 'vector_embeddings'

    id = Column(Integer, primary_key=True, autoincrement=True)
    doc_id = Column(Integer, nullable=False)  # 关联文档ID
    chunk_index = Column(Integer, nullable=False)  # 文本块索引
    chunk_text = Column(Text, nullable=False)  # 原始文本
    embedding = Column(JSON, nullable=False)  # 向量数据
    created_at = Column(DateTime, default=datetime.utcnow)

    def __repr__(self):
        return f"<VectorEmbedding(doc_id={self.doc_id}, chunk_index={self.chunk_index})>"


class RecallHistory(Base):
    """召回历史，存储每次召回的查询与命中摘要"""
    __tablename__ = 'recall_history'

    id = Column(Integer, primary_key=True, autoincrement=True)
    kb_id = Column(Integer, ForeignKey("knowledge_bases.id", ondelete="CASCADE"), nullable=False)
    query_text = Column(Text, nullable=False)
    hits = Column(JSON, nullable=False)  # [{doc_id, doc_name, chunk_index, chunk_text, score}]
    created_at = Column(DateTime, default=datetime.utcnow)


# 用户相关模型
class User(Base):
    """用户表"""
    __tablename__ = "users"

    user_id = Column(Integer, primary_key=True, autoincrement=True)
    username = Column(String(50), nullable=False, unique=True)
    email = Column(String(100), nullable=False, unique=True)
    password_hash = Column(String(255), nullable=False)
    role_id = Column(Integer, ForeignKey("roles.role_id"), nullable=False, default=2)
    is_active = Column(Boolean, nullable=False, default=True)
    created_time = Column(DateTime, nullable=False, default=datetime.utcnow)
    updated_time = Column(DateTime, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)
    last_login_time = Column(DateTime, nullable=True)

    # 关联关系
    role = relationship("Role", back_populates="users")
    conversations = relationship("Conversation", back_populates="user")
    system_prompts = relationship("SystemPrompt", back_populates="user")
    knowledge_bases = relationship("KnowledgeBase", back_populates="user")
    tokens = relationship("UserToken", back_populates="user", cascade="all, delete-orphan")


class Role(Base):
    """角色表"""
    __tablename__ = "roles"

    role_id = Column(Integer, primary_key=True, autoincrement=True)
    role_name = Column(String(50), nullable=False, unique=True)
    role_description = Column(String(255), nullable=True)
    created_time = Column(DateTime, nullable=False, default=datetime.utcnow)

    # 关联关系
    users = relationship("User", back_populates="role")
    permissions = relationship("Permission", secondary="role_permissions", back_populates="roles")


class Permission(Base):
    """权限表"""
    __tablename__ = "permissions"

    permission_id = Column(Integer, primary_key=True, autoincrement=True)
    permission_name = Column(String(50), nullable=False, unique=True)
    permission_description = Column(String(255), nullable=True)
    created_time = Column(DateTime, nullable=False, default=datetime.utcnow)

    # 关联关系
    roles = relationship("Role", secondary="role_permissions", back_populates="permissions")


class RolePermission(Base):
    """角色权限关联表"""
    __tablename__ = "role_permissions"

    role_id = Column(Integer, ForeignKey("roles.role_id", ondelete="CASCADE"), primary_key=True)
    permission_id = Column(Integer, ForeignKey("permissions.permission_id", ondelete="CASCADE"), primary_key=True)


class UserToken(Base):
    """用户token表"""
    __tablename__ = "user_tokens"

    token_id = Column(Integer, primary_key=True, autoincrement=True)
    user_id = Column(Integer, ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False)
    token_hash = Column(String(255), nullable=False)
    expires_at = Column(DateTime, nullable=False)
    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
    is_active = Column(Boolean, nullable=False, default=True)

    # 关联关系
    user = relationship("User", back_populates="tokens")
